from model import Exam, Question, UnsupportedQuestion, ExamCollection

from bs4 import BeautifulSoup, Tag
from loguru import logger
from typing import Optional, Dict
from os import path

import requests


provinces = ['国考', '浙江', '山东', '江苏', '广东']
question_types = set([
    '政治理论', '常识判断', '言语理解与表达',
    '数量关系', '判断推理', '资料分析',
    '数学运算', '类比推理', '图形推理', '逻辑判断', '科学推理'
])


def collect_question(question_type: str, row: Tag) -> Question:
    question_no = int(row.find('div', attrs={'class': 'left'}).get_text())
    content_div: Tag = row.find('div', attrs={'class': 'right'})
    for uitem in content_div.findAll('u'):
        if uitem.getText().strip() == '':
            uitem.replaceWith("________")

    question_content = row.find('div', attrs={'class': 'right'}).get_text()
    question_content = question_content.replace('A、', ' 选项列表：A、')
    if row.find('img') is not None:
        raise UnsupportedQuestion(f"question {question_no} contains unsupported image information")

    return Question(id=question_no, type=question_type, content=question_content)


def collect_exam(exam_id: str) -> Exam:
    logger.info(f"collect exam information at {exam_id}")
    exam = Exam(id=exam_id, title='')
    paper_resp = requests.get(f"https://www.gkzenti.cn/paper/{exam_id}")
    answers_resp = requests.get(f"https://www.gkzenti.cn/answer/{exam_id}")

    if paper_resp.status_code != 200:
        logger.critical(f"failed to get questions for exam {exam_id}")
        return
    if answers_resp.status_code != 200:
        logger.critical(f"failed to get answers for exam {exam_id}")
        return

    # 首先获取所有答案并记录下来
    answers: Dict[int, str] = dict()
    soup = BeautifulSoup(answers_resp.text, 'html.parser')
    container: Tag = soup.find('div', attrs={'id': 'printcontent'})
    for tag in container.find_all('div', attrs={'class': 'col-xs-1-5'}):
        qid, answer = tag.getText().split('、')
        answers[int(qid)] = answer

    soup = BeautifulSoup(paper_resp.text, 'html.parser')
    container = soup.find('div', attrs={'id': 'printcontent'})
    curr_question_type = ""
    curr_section_title: Optional[str] = None
    curr_section_header: Optional[str] = None

    for row in container.find_all('div', attrs={'class': 'row'}):
        assert isinstance(row, Tag)
        subtitle = row.find('div', attrs={'class': 'subtitle'})
        sub2title = row.find('div', attrs={'class': 'sub2title'})

        if subtitle is not None:
            # 这是一个题型的开始
            subtitle_text = subtitle.getText()
            for ty in question_types:
                if ty in subtitle_text:
                    question_type = ty
                    break
            else:
                logger.critical(f"unknown question type from: {subtitle_text}")
                logger.info(f"{row}")
                raise SystemExit(1)

            curr_question_type = question_type
        elif sub2title is not None:
            section_title = sub2title.getText()
            curr_section_title = section_title
            if row.find('img') is not None:
                logger.warning(f"section {section_title} contains unsupported image information")
                curr_section_header = None
            else:
                curr_section_header = row.findChildren('div')[1].getText()
        else:
            if curr_section_title is not None and curr_section_header is None:
                # 我们当前处在某一个题目的子章节中, e.g. 阅读理解，但是章节信息中可能由于
                # 包含多模态信息导致我们无法解析
                continue

            try:
                question = collect_question(curr_question_type, row)
                question.answer = answers[question.id]
                if curr_section_header is not None:
                    question.content = f"背景资料: {curr_section_header}。请基于背景资料回答问题：{question.content}"

                exam.questions.append(question)
            except UnsupportedQuestion as ex:
                logger.warning(f"{ex}")
            except Exception as ex:
                logger.warning(f"{ex}: {row}")
                breakpoint()

    return exam


def main() -> None:
    collection = ExamCollection()
    for province in provinces:
        resp = requests.get(f"https://www.gkzenti.cn/api/json?cls=行测&province={province}")
        for item in resp.json():
            title = item['Title']
            if title[:4] in ('2024', '2025'):
                exam_id = item['No'].split('/')[-1]
                exam = collect_exam(exam_id)
                exam.title = title
                collection.exams.append(exam)

    logger.info(f"{len(collection.exams)} examinations collected in total")
    script_path = path.dirname(path.realpath(__file__))
    with open(path.join(script_path, "data.json"), "w+") as fdata:
        fdata.write(collection.model_dump_json(indent=2))


if __name__ == "__main__":
    main()
